# coding=utf-8
import os
from Crypto import Random
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher
from Crypto.Signature import PKCS1_v1_5 as PKCS1_signature
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex

import base64
# RSA加密传输数据
# AES加密传输公钥

class AESCipher:
    """
        RSA加密、解密、签名、验签工具类
    """
    def __init__(self, key='5YyX5p6X6Z+z5bqm'):
        self.mode = AES.MODE_OFB
        self.aeskey = key

    def encrypt(self, text):
        """
           加密方法
           :param text: 需要加密的明文
           :return: 加密后的密文
       """
        if os.path.exists('./secretkey/public_a.rsa'):
            with open('./secretkey/public_a.rsa') as f:
                key = f.read()
                pub_key = RSA.importKey(str(key))
                cipher = PKCS1_cipher.new(pub_key)
                rsa_text = base64.b64encode(cipher.encrypt(bytes(text.encode("utf8"))))
                # return self.encryptByAes(rsa_text.decode('utf-8'))
                return rsa_text.decode('utf-8')
        else:
            self.generateKeys()
            self.encrypt(text)

    def decrypt(self, Ciphertext):
        """
           解密方法
           :param Ciphertext: 加密后的密文
           :return: 解密后的明文
       """
        # Ciphertext = self.decryptByAes(Ciphertext)
        if os.path.exists('./secretkey/private_a.rsa'):
            # 使用私钥对内容进行rsa解密
            with open('./secretkey/private_a.rsa') as f:
                key = f.read()
                pri_key = RSA.importKey(key)
                cipher = PKCS1_cipher.new(pri_key)
                back_text = cipher.decrypt(base64.b64decode(Ciphertext), 0)
                return back_text.decode('utf-8')
        else:
            raise Exception('private key does not exist ')

    def generateKeys(self):
        random_generator = Random.new().read
        rsa = RSA.generate(2048, random_generator)
        private_key = rsa.exportKey()
        with open("./secretkey/private_a.rsa", 'wb') as f:
            f.write(private_key)
        public_key = rsa.publickey().exportKey()
        with open("./secretkey/public_a.rsa", 'wb') as f:
            f.write(public_key)

    def signText(self, text):

        # 使用私钥生成签名
        with open('./secretkey/private_a.rsa') as f:
            key = f.read()
            pri_key = RSA.importKey(key)
            signer = PKCS1_signature.new(pri_key)
            digest = SHA.new()
            digest.update(text.encode("utf8"))
            sign = signer.sign(digest)
            signature = base64.b64encode(sign)
            return signature.decode('utf-8')

    def verify(self, text, signature):
        # 使用公钥验证签名
        with open('./secretkey/public_a.rsa') as f:
            key = f.read()
            pub_key = RSA.importKey(key)
            verifier = PKCS1_signature.new(pub_key)
            digest = SHA.new()
            digest.update(text.encode("utf8"))
            return verifier.verify(digest, base64.b64decode(signature))

    def getPublicKey(self):
        '''
        :return: 返回公钥
        '''
        if os.path.exists('./secretkey/public_a.rsa'):
            with open('./secretkey/public_a.rsa') as f:
                key = f.read()
                pub_key = RSA.importKey(str(key)).exportKey()
                return pub_key
        else:
            self.generateKeys()
            self.getPublicKey()

    def encryptByAes(self, message):
        cryptor = AES.new(self.aeskey.encode('utf-8'), self.mode, b'0000000000000000')
        length = 16
        count = len(message)
        if count % length != 0:
            add = length - (count % length)
        else:
            add = 0
        message = message + ('\0' * add)
        ciphertext = cryptor.encrypt(message.encode('utf-8'))
        result = b2a_hex(ciphertext)
        return result.decode('utf-8')

    def decryptByAes(self, result):
        cryptor = AES.new(self.aeskey.encode('utf-8'), self.mode, b'0000000000000000')
        plain_text = cryptor.decrypt(a2b_hex(result))
        return plain_text.decode('utf-8').rstrip('\0')


# if __name__ == '__main__':
#     aesdemo = AESCipher()
#     msg = '你好'
#     encrypt = aesdemo.encryptByAes(msg)
#     print(encrypt)
#     print('*' * 10)
#     decrypt = aesdemo.decryptByAes(encrypt)
#     print(decrypt)
#     sign = aesdemo.signText(encrypt)
#     print(sign)
#     verify1 = aesdemo.verify(encrypt, sign)
#     print(verify1)
    # pubKey = aesdemo.getPublicKey()
    # print(pubKey)



